#ifndef TENSORRT_DETECTOR_H
#define TENSORRT_DETECTOR_H

#include <iostream>
#include <vector>
#include <memory>
#include <opencv2/opencv.hpp>
#include <NvInfer.h>

// Forward declarations
namespace nvinfer1 {
    class IRuntime;
    class ICudaEngine;
    class IExecutionContext;
}

// Destroy TensorRT objects
struct TRTDestroy {
    template <class T>
    void operator()(T* obj) const {
        if (obj)
            obj->destroy();
    }
};

template <class T>
using TRTUniquePtr = std::unique_ptr<T, TRTDestroy>;

// TensorRT Detector Class
class TensorRTDetector {
public:
    TensorRTDetector(const std::string& model_path, const std::string& plugin_path);
    ~TensorRTDetector();
    bool initialize();
    std::vector<float> inferenceFromFile(const std::string& input_file_path);
    std::vector<float> inferenceFromImage(const std::string& image_path);
    std::pair<std::vector<float>, long long> inferenceFromImageWithTime(const std::string& image_path);
    bool visualizeResults(const std::vector<float>& output_data, const std::string& output_path);
    std::vector<float> applyNMS(const std::vector<float>& bboxes, float conf_threshold = 0.3f, float iou_threshold = 0.8f);

private:
    bool loadInputData(const std::string& file_path, std::vector<float>& data);
    bool preprocessImage(const std::string& image_path, std::vector<float>& processed_data);
    std::vector<float> performRealInference(const std::vector<float>& input_data);
    std::vector<float> performMockInference(const std::vector<float>& input_data);
    float calculateIoU(const std::vector<float>& bbox1, const std::vector<float>& bbox2);
    size_t getOutputSize();
    size_t getInputElementCount();
    size_t getOutputElementCount();

    std::string model_path_;
    std::string plugin_path_;
    bool initialized_;
    
    // TensorRT objects
    TRTUniquePtr<nvinfer1::IRuntime> runtime_;
    TRTUniquePtr<nvinfer1::ICudaEngine> engine_;
    TRTUniquePtr<nvinfer1::IExecutionContext> context_;
    
    // Input/Output dimensions
    nvinfer1::Dims input_dims_;
    nvinfer1::Dims output_dims_;
};

#endif // TENSORRT_DETECTOR_H
